package com.novitechie;

import com.janetfilter.core.plugin.MyTransformer;
import jdk.internal.org.objectweb.asm.ClassReader;
import jdk.internal.org.objectweb.asm.ClassWriter;
import jdk.internal.org.objectweb.asm.tree.*;

import static jdk.internal.org.objectweb.asm.Opcodes.*;

/**
 * @author YeloChick
 */
public class SystemTransformer implements MyTransformer {

    @Override
    public String getHookClassName() {
        return "java/lang/System";
    }

    @Override
    public byte[] transform(String className, byte[] classBytes, int order) throws Exception {
        ClassReader reader = new ClassReader(classBytes);
        ClassNode node = new ClassNode(ASM5);
        reader.accept(node, 0);
        for (MethodNode m : node.methods) {
            if ("getenv".equals(m.name) && m.desc.equals("(Ljava/lang/String;)Ljava/lang/String;")) {
                InsnList list = new InsnList();
                LabelNode L0 = new LabelNode();
                list.add(new VarInsnNode(ALOAD, 0));
                list.add(new MethodInsnNode(INVOKESTATIC, "com/novitechie/rules/SystemRule", "checkEnv", "(Ljava/lang/String;)Z", false));
                list.add(new JumpInsnNode(IFEQ, L0));
                list.add(new InsnNode(ACONST_NULL));
                list.add(new InsnNode(ARETURN));
                list.add(L0);
                m.instructions.insert(list);
            } else if ("getProperty".equals(m.name) && (m.desc.equals("(Ljava/lang/String;)Ljava/lang/String;") || m.desc.equals("(Ljava/lang/String;Ljava/lang/String;)Ljava/lang/String;"))) {
                InsnList list = new InsnList();
                LabelNode L0 = new LabelNode();
                list.add(new VarInsnNode(ALOAD, 0));
                list.add(new MethodInsnNode(INVOKESTATIC, "com/novitechie/rules/SystemRule", "checkProperty", "(Ljava/lang/String;)Z", false));
                list.add(new JumpInsnNode(IFEQ, L0));
                list.add(new InsnNode(ACONST_NULL));
                list.add(new InsnNode(ARETURN));
                list.add(L0);
                m.instructions.insert(list);
            }
        }
        ClassWriter writer = new ClassWriter(ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS);
        node.accept(writer);
        return writer.toByteArray();
    }
}
